package com.zengde.ai_sandbox.handler;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.zengde.ai_sandbox.manager.BigModelManager;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
import org.springframework.context.annotation.Profile;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
@ConditionalOnWebApplication(type = ConditionalOnWebApplication.Type.SERVLET)
@Profile("!test")
@RequiredArgsConstructor
@Slf4j
public class RawWebSocketHandler extends TextWebSocketHandler {
    
    private final BigModelManager bigModelManager;
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    // 存储会话关系: sessionId -> WebSocketSession
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        log.info("WebSocket连接已建立: {}, 远程地址: {}", session.getId(), session.getRemoteAddress());
        sessions.put(session.getId(), session);
        
        // 发送连接成功消息
        try {
            session.sendMessage(new TextMessage("{\"type\":\"CONNECTED\",\"message\":\"连接成功\"}"));
            log.info("发送连接成功消息完成");
        } catch (IOException e) {
            log.error("发送连接成功消息失败", e);
        }
    }
    
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        log.error("WebSocket传输错误: sessionId={}, 错误: {}", session.getId(), exception.getMessage(), exception);
        try {
            session.close(CloseStatus.SERVER_ERROR.withReason("传输错误: " + exception.getMessage()));
        } catch (IOException e) {
            log.error("关闭WebSocket会话失败", e);
        } finally {
            sessions.remove(session.getId());
        }
    }
    
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        try {
            String payload = message.getPayload();
            log.info("收到WebSocket消息: {}, sessionId: {}", payload, session.getId());
            
            // 解析JSON消息
            JsonNode jsonNode = objectMapper.readTree(payload);
            String action = jsonNode.path("action").asText();
            
            log.info("处理action: {}", action);
            
            // 根据action类型处理不同的请求
            switch (action) {
                case "createSession":
                    handleCreateSession(session, jsonNode);
                    break;
                case "sendMessage":
                    handleSendMessage(session, jsonNode);
                    break;
                case "test":
                    handleTest(session, jsonNode);
                    break;
                default:
                    log.warn("未知的action类型: {}", action);
                    sendErrorMessage(session, "未知的action类型: " + action);
            }
        } catch (Exception e) {
            log.error("处理WebSocket消息异常: {}", e.getMessage(), e);
            try {
                String errorMsg = "{\"type\":\"ERROR\",\"message\":\"" + escapeJson(e.getMessage()) + "\"}";
                log.info("发送错误响应: {}", errorMsg);
                session.sendMessage(new TextMessage(errorMsg));
            } catch (IOException ex) {
                log.error("发送错误消息失败", ex);
            }
        }
    }
    
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        log.info("WebSocket连接已关闭: {}, 状态: {}", session.getId(), status);
        sessions.remove(session.getId());
    }
    
    /**
     * 处理创建会话请求
     */
    private void handleCreateSession(WebSocketSession session, JsonNode jsonNode) throws IOException {
        String customerId = jsonNode.path("customerId").asText();
        String modelServiceName = jsonNode.path("modelServiceName").asText("doubaoService");
        
        log.info("创建会话请求: customerId={}, modelServiceName={}", customerId, modelServiceName);
        
        try {
            // 创建会话
            String sessionId = bigModelManager.createSession(customerId, modelServiceName);
            log.info("会话创建成功: {}", sessionId);
            
            // 获取机器人首次发言
            String firstMessage = bigModelManager.robotFirstMessage(sessionId);
            log.info("机器人首次发言: {}", firstMessage);
            
            // 构建响应消息
            String response = String.format(
                "{\"type\":\"SESSION_CREATED\",\"sessionId\":\"%s\",\"firstMessage\":\"%s\"}", 
                sessionId, 
                escapeJson(firstMessage)
            );
            
            // 发送响应
            log.info("发送会话创建响应: {}", response);
            session.sendMessage(new TextMessage(response));
            log.info("会话创建响应已发送");
        } catch (Exception e) {
            log.error("创建会话失败: {}", e.getMessage(), e);
            sendErrorMessage(session, "创建会话失败: " + e.getMessage());
        }
    }
    
    /**
     * 处理发送消息请求
     */
    private void handleSendMessage(WebSocketSession session, JsonNode jsonNode) throws IOException {
        String sessionId = jsonNode.path("sessionId").asText();
        String content = jsonNode.path("content").asText();
        String customerId = jsonNode.path("customerId").asText();
        String modelType = jsonNode.path("modelType").asText("normal");
        String responseType = jsonNode.path("responseType").asText("stream");
        
        log.info("发送消息请求: sessionId={}, content={}, modelType={}, responseType={}", 
                sessionId, content, modelType, responseType);
        
        try {
            // 根据不同的模型类型和响应类型选择调用方式
            if ("think".equals(modelType)) {
                if ("stream".equals(responseType)) {
                    // 深度思考、流式调用
                    handleStreamResponse(session, bigModelManager.ts(sessionId, content), sessionId);
                } else {
                    // 深度思考、非流式调用
                    String response = bigModelManager.tns(sessionId, content);
                    sendFullResponse(session, sessionId, response);
                }
            } else {
                if ("stream".equals(responseType)) {
                    // 非深度思考、流式调用
                    handleStreamResponse(session, bigModelManager.nts(sessionId, content), sessionId);
                } else {
                    // 非深度思考、非流式调用
                    String response = bigModelManager.ntns(sessionId, content);
                    sendFullResponse(session, sessionId, response);
                }
            }
        } catch (Exception e) {
            log.error("发送消息失败", e);
            sendErrorMessage(session, "发送消息失败: " + e.getMessage());
        }
    }
    
    /**
     * 处理测试请求
     */
    private void handleTest(WebSocketSession session, JsonNode jsonNode) throws IOException {
        log.info("处理测试请求: {}", jsonNode);
        String response = String.format(
            "{\"type\":\"TEST_RESPONSE\",\"message\":\"这是一个测试响应\",\"timestamp\":%d}",
            System.currentTimeMillis()
        );
        log.info("发送测试响应: {}", response);
        session.sendMessage(new TextMessage(response));
        log.info("测试响应已发送");
    }
    
    /**
     * 处理流式响应
     */
    private void handleStreamResponse(WebSocketSession session, Flux<String> responseFlux, String sessionId) {
        responseFlux.subscribe(
            chunk -> {
                try {
                    String response = String.format(
                        "{\"type\":\"RESPONSE_CHUNK\",\"sessionId\":\"%s\",\"content\":\"%s\"}",
                        sessionId, 
                        escapeJson(chunk)
                    );
                    session.sendMessage(new TextMessage(response));
                } catch (IOException e) {
                    log.error("发送流式响应片段失败", e);
                }
            },
            error -> {
                log.error("流式响应发生错误", error);
                try {
                    sendErrorMessage(session, "流式响应错误: " + error.getMessage());
                } catch (IOException e) {
                    log.error("发送错误消息失败", e);
                }
            },
            () -> {
                try {
                    String response = String.format(
                        "{\"type\":\"RESPONSE_COMPLETE\",\"sessionId\":\"%s\"}",
                        sessionId
                    );
                    session.sendMessage(new TextMessage(response));
                    log.info("流式响应完成消息已发送");
                } catch (IOException e) {
                    log.error("发送流式响应完成消息失败", e);
                }
            }
        );
    }
    
    /**
     * 发送完整响应
     */
    private void sendFullResponse(WebSocketSession session, String sessionId, String content) throws IOException {
        String response = String.format(
            "{\"type\":\"FULL_RESPONSE\",\"sessionId\":\"%s\",\"content\":\"%s\"}",
            sessionId, 
            escapeJson(content)
        );
        session.sendMessage(new TextMessage(response));
        log.info("完整响应已发送");
    }
    
    /**
     * 发送错误消息
     */
    private void sendErrorMessage(WebSocketSession session, String errorMessage) throws IOException {
        String response = String.format(
            "{\"type\":\"ERROR\",\"message\":\"%s\"}",
            escapeJson(errorMessage)
        );
        log.info("发送错误消息: {}", response);
        session.sendMessage(new TextMessage(response));
        log.info("错误消息已发送: {}", errorMessage);
    }
    
    /**
     * 转义JSON字符串
     */
    private String escapeJson(String input) {
        if (input == null) {
            return "";
        }
        return input.replace("\\", "\\\\")
                .replace("\"", "\\\"")
                .replace("\n", "\\n")
                .replace("\r", "\\r")
                .replace("\t", "\\t");
    }
} 